package xjay.ai.emotionalsupport.rag;



import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Slf4j
@Data
public abstract class RagService {

    private VectorStore vectorStore;
    private  ResourcePatternResolver resolver;
    private String pdfResourcePath;
    private String markdownResourcePath;

    protected List<Document> documents = new ArrayList<>();
    protected List<Resource> resources = new ArrayList<>();

    public void loadMarkdown() throws IOException {
        log.info("load markdown"); //有多少个文档就打印多少次
        Resource[] markdownResources = resolver.getResources(markdownResourcePath);
        this.resources.addAll(Arrays.asList(markdownResources));
        // 创建文本分片器
        TokenTextSplitter textSplitter = new TokenTextSplitter(
                500,  // 每个chunk的token数量（减小长度以避免API限制）
                50,   // 重叠token数量
                10,   // 最小chunk长度
                1000, // 最大chunk长度（减小以避免API限制）
                true  // 保留分隔符
        );
        
        for (Resource resource : markdownResources) {
            String filename = resource.getFilename();
            MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                    .withHorizontalRuleCreateDocument(true)
                    .withIncludeCodeBlock(false)
                    .withIncludeBlockquote(false)
                    .withAdditionalMetadata("filename", filename)
                    .build();

            MarkdownDocumentReader reader = new MarkdownDocumentReader(resource,config);
            List<Document> rawDocuments = reader.get();
            
            // 对每个文档进行分片处理
            for (Document doc : rawDocuments) {
                List<Document> chunks = textSplitter.split(doc);
                documents.addAll(chunks);
            }
        }

    }

    public void loadPDF() throws IOException {
        log.info("load PDF"); //有多少pdf打印多少次
        Resource[] pdfResources = resolver.getResources(pdfResourcePath);
        //数组转换成列表：Arrays.asList(数组)
        this.resources.addAll(Arrays.asList(pdfResources));
        for (Resource resource : pdfResources) {
            String filename = resource.getFilename();
            PdfDocumentReaderConfig config = PdfDocumentReaderConfig.builder()
                    .withPageTopMargin(0)
                    .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                            .withNumberOfTopTextLinesToDelete(0)
                            .build())
                    .withPagesPerDocument(1)
                    .build();

            PagePdfDocumentReader reader = new PagePdfDocumentReader(resource,config);
            documents.addAll(reader.get());
        }
    }

    // 删除特定文档的方法
    public void deleteDocument(String filename) {
        // 按文件名删除
        getVectorStore().delete("filename == '" + filename + "'"); //metadata就是一个Map
    }

    // 删除多个文档ID
    public void deleteDocuments(List<String> documentIds) {
        getVectorStore().delete(documentIds);
    }

    // 重新加载文档（删除旧的，添加新的）
    public void refreshDocuments() throws IOException {
        // 删除所有占卜相关文档
        getVectorStore().delete("filename LIKE '%.md'");

        // 重新加载
        loadMarkdown();
        rag();
    }

    public abstract VectorStore rag() throws IOException;


}
